"use client";
import { useState, useRef, useEffect } from "react";
import {
  Select,
  SelectContent,
  SelectItem,
  SelectTrigger,
  SelectValue,
} from "@/components/ui/select";
import { ScrollArea } from "@/components/ui/scroll-area";
import { Button } from "@/components/ui/button";
import { ReactFlowProvider } from "@xyflow/react";
import ValidationFlow from "@/components/agent/validation/validation-flow-v2";
import { TestDataGenerator } from "@/lib/utils";
import { config } from "@/config";
import { CheckCircle, Circle, Terminal, Inbox } from "lucide-react"; // 合并 Inbox 导入
import {
  Dialog,
  DialogContent,
  DialogHeader,
  DialogTitle,
  DialogTrigger,
  DialogDescription,
} from "@/components/ui/dialog";
import { useSSE } from "@/hooks/sse-connector";
import {
  Card,
  CardContent,
  CardDescription,
  CardHeader,
  CardTitle,
} from "@/components/ui/card";
import { ModalType, TaskType, Algorithm } from "@/lib/types";
import { ModalMap } from "@/lib/const";

export default function ValidationPage() {
  const [algorithmMap, setAlgorithmMap] = useState<Record<string, Algorithm[]>>(
    {}
  );
  const [selectedTask, setSelectedTask] = useState<TaskType | "">("");
  const [selectedAlgorithm, setSelectedAlgorithm] = useState<string>("");
  const [selectedConfig, setSelectedConfig] = useState<string>("");
  const [selectedWeight, setSelectedWeight] = useState<string>("");
  const [selectedModal, setSelectedModal] = useState<ModalType | "">("");

  useEffect(() => {
    fetch(
      `${config.AI_SERVICE_URL}${config.AI_SERVICE_URL_PREFIX}/get_yaml_files`
    )
      .then((response) => response.json())
      .then((data) => {
        if (data.error) return;
        const algorithms: Algorithm[] = data.algorithms || [];
        const groupedData =
          algorithms.reduce(
            (acc: Record<TaskType, Algorithm[]>, item: Algorithm) => {
              if (!acc[item.task]) {
                acc[item.task] = [];
              }
              acc[item.task].push(item);
              return acc;
            },
            {} as Record<TaskType, Algorithm[]>
          ) ?? [];
        setAlgorithmMap(groupedData);
      });
  }, []);

  useEffect(() => {
    setSelectedAlgorithm("");
    setSelectedConfig("");
    setSelectedWeight("");
    setSelectedModal("");
  }, [selectedTask]);

  useEffect(() => {
    setSelectedConfig("");
    setSelectedWeight("");
    setSelectedModal("");
  }, [selectedAlgorithm]);

  const { messages, sseData, connect } = useSSE();
  const [currentDataIndex, setCurrentDataIndex] = useState<number>(0);
  // 日志自动滚动到底部
  const logRef = useRef<HTMLDivElement>(null);

  // Helper function to switch image group
  const handleSwitchImageGroup = (direction: "up" | "down") => {
    setCurrentDataIndex((prevIndex) => {
      if (direction === "up") {
        return Math.max(0, prevIndex - 1);
      }
      return Math.min(sseData.length - 1, prevIndex + 1);
    });
  };

  useEffect(() => {
    if (logRef.current) {
      logRef.current.scrollTop = logRef.current.scrollHeight;
    }
  }, [messages]);

  useEffect(() => {
    setCurrentDataIndex(sseData.length - 1);
  }, [sseData.length]);

  const [isPreviewOpen, setIsPreviewOpen] = useState(false);
  const [previewContent, setPreviewContent] = useState<string>("");

  const handlePreviewConfig = async () => {
    if (!selectedConfig) return;
    try {
      const response = await fetch(
        `${config.AI_SERVICE_URL}${config.AI_SERVICE_URL_PREFIX}/get_config?config_path=${selectedConfig}`
      );
      const data = await response.json();
      if (data.error) {
        setPreviewContent(`Error: ${data.error}`);
      } else {
        setPreviewContent(JSON.stringify(data, null, 2)); // Ensure data.content is used and formatted
      }
      setIsPreviewOpen(true);
    } catch (error) {
      setPreviewContent(`Error fetching config: ${error}`);
      setIsPreviewOpen(true);
    }
  };

  const handleStartValidation = () => {
    connect(
      `${config.AI_SERVICE_URL}${config.AI_SERVICE_URL_PREFIX}/rl_test?env_id=${selectedTask}&algorithm_name=${selectedAlgorithm}&validation_config=${selectedConfig}&model_path=${selectedWeight}&modal=${selectedModal}`,
      {
        generator: TestDataGenerator,
      }
    );
  };

  return (
    <div className="flex-1 h-full m-4 flex gap-4">
      <ScrollArea className="h-full">
        <div className="space-y-4">
          <Card className="w-[350px]">
            <CardHeader>
              <CardTitle>
                <div className="flex items-center">
                  {selectedTask ? (
                    <CheckCircle className="w-5 h-5 mr-2 text-green-500" />
                  ) : (
                    <Circle className="w-5 h-5 mr-2 text-gray-400" />
                  )}
                  <span className="font-bold text-base">STEP1: 选择任务</span>
                </div>
              </CardTitle>
              <CardDescription>
                选择要推理的任务,包括预处理-Preprocessing,分类-classification,检测-Detection,跟踪-Tracking,匹配搜索-Search
              </CardDescription>
            </CardHeader>
            <CardContent>
              <Select
                value={selectedTask}
                onValueChange={(value: TaskType) => {
                  setSelectedTask(value);
                }}
              >
                <SelectTrigger className="w-full">
                  <SelectValue placeholder="请选择任务" />
                </SelectTrigger>
                <SelectContent>
                  {Object.keys(algorithmMap).map((type: string) => (
                    <SelectItem key={type} value={type}>
                      {type}
                    </SelectItem>
                  ))}
                </SelectContent>
              </Select>
            </CardContent>
          </Card>

          <Card className="w-[350px]">
            <CardHeader>
              <CardTitle>
                <div className="flex items-center">
                  {selectedAlgorithm ? (
                    <CheckCircle className="w-5 h-5 mr-2 text-green-500" />
                  ) : (
                    <Circle className="w-5 h-5 mr-2 text-gray-400" />
                  )}
                  <span className="font-bold text-base">STEP2: 选择算法</span>
                </div>
              </CardTitle>
              <CardDescription>选择要测试的强化学习算法</CardDescription>
            </CardHeader>
            <CardContent>
              <Select
                disabled={!selectedTask}
                value={selectedAlgorithm}
                onValueChange={(value) => {
                  setSelectedAlgorithm(value);
                }}
              >
                <SelectTrigger className="w-full">
                  <SelectValue placeholder="请选择算法" />
                </SelectTrigger>
                <SelectContent>
                  {algorithmMap[selectedTask]?.map((alg) => (
                    <SelectItem key={alg.name} value={alg.name}>
                      {alg.name}
                    </SelectItem>
                  ))}
                </SelectContent>
              </Select>
            </CardContent>
          </Card>
          <Card className="w-[350px]">
            <CardHeader>
              <CardTitle>
                <div className="flex items-center">
                  {selectedConfig && selectedWeight && selectedModal ? (
                    <CheckCircle className="w-5 h-5 mr-2 text-green-500" />
                  ) : (
                    <Circle className="w-5 h-5 mr-2 text-gray-400" />
                  )}
                  <span className="font-bold text-base">
                    STEP3: 设置配置信息
                  </span>
                </div>
              </CardTitle>
              <CardDescription>
                同一个算法可能有多个配置文件和权重文件，还有不同的模态，请进行选择
              </CardDescription>
            </CardHeader>
            <CardContent>
              <div className="flex gap-2 w-full">
                <Select
                  disabled={!selectedAlgorithm}
                  value={selectedConfig}
                  onValueChange={setSelectedConfig}
                >
                  <SelectTrigger className="w-3/4">
                    <SelectValue placeholder="请选择配置文件" />
                  </SelectTrigger>
                  <SelectContent>
                    {algorithmMap[selectedTask]
                      ?.find((alg) => alg.name === selectedAlgorithm)
                      ?.validation_config?.map((cfg: string) => (
                        <SelectItem key={cfg} value={cfg}>
                          {cfg}
                        </SelectItem>
                      ))}
                  </SelectContent>
                </Select>
                <Dialog open={isPreviewOpen} onOpenChange={setIsPreviewOpen}>
                  <DialogTrigger asChild>
                    <Button
                      className="w-1/4"
                      variant="secondary"
                      disabled={!Boolean(selectedConfig)}
                      onClick={handlePreviewConfig} // Call fetch on click, dialog opens via state
                    >
                      预览
                    </Button>
                  </DialogTrigger>
                  <DialogContent className="sm:max-w-[600px]">
                    <DialogHeader>
                      <DialogTitle>配置文件预览</DialogTitle>
                      <DialogDescription>
                        以下是选定配置文件的内容。
                      </DialogDescription>
                    </DialogHeader>
                    <ScrollArea className="h-[600px] w-full rounded-md border p-4 bg-muted/40">
                      <pre className="text-sm whitespace-pre-wrap break-all">
                        {previewContent}
                      </pre>
                    </ScrollArea>
                  </DialogContent>
                </Dialog>
              </div>
              <Select
                disabled={!selectedAlgorithm}
                value={selectedWeight}
                onValueChange={setSelectedWeight}
              >
                <SelectTrigger className="w-full mt-2">
                  <SelectValue placeholder="请选择权重文件" />
                </SelectTrigger>
                <SelectContent>
                  {algorithmMap[selectedTask]
                    ?.find((alg) => alg.name === selectedAlgorithm)
                    ?.weights?.map((weight: string) => (
                      <SelectItem key={weight} value={weight}>
                        {weight}
                      </SelectItem>
                    ))}
                </SelectContent>
              </Select>
              <Select
                disabled={!selectedAlgorithm}
                value={selectedModal}
                onValueChange={(value: ModalType) => setSelectedModal(value)}
              >
                <SelectTrigger className="w-full mt-2">
                  <SelectValue placeholder="请选择模态" />
                </SelectTrigger>
                <SelectContent>
                  {selectedTask &&
                    ModalMap[selectedTask]?.map((item) => (
                      <SelectItem key={item} value={item}>
                        {item}
                      </SelectItem>
                    ))}
                </SelectContent>
              </Select>
            </CardContent>
          </Card>
          <Button
            disabled={!selectedConfig || !selectedModal || !selectedWeight}
            onClick={handleStartValidation}
          >
            开始推理
          </Button>
        </div>
      </ScrollArea>
      <div className="flex-1 border rounded-md flex flex-col gap-2 p-4">
        <div className="text-sm font-medium flex gap-2">
          <Terminal className="h-5 w-5" />
          <span>推理日志</span>
        </div>
        <div
          ref={logRef}
          className="bg-neutral-900/95 rounded-md p-4 h-[180px] overflow-y-auto border font-mono text-xs text-green-400"
        >
          {messages.length === 0 && "推理未开始"}
          {messages.map((log: string, index: number) => (
            <div key={index} className="mb-1 whitespace-pre-wrap break-all">
              <span className="text-green-600">$</span> {log}
            </div>
          ))}
        </div>
        <div className="flex-1">
          {sseData.length === 0 ? (
            <div className="flex flex-col items-center justify-center h-full text-muted-foreground">
              <Inbox className="w-16 h-16 mb-4" />
              <p className="text-lg font-medium">暂无推理数据</p>
              <p className="text-sm">请先开始推理以查看可视化结果。</p>
            </div>
          ) : (
            <ReactFlowProvider>
              <ValidationFlow
                rlName={selectedAlgorithm}
                cvModels={[sseData[currentDataIndex]?.chosenModel as string]}
                inputImages={
                  (sseData[currentDataIndex]?.inputImages as string[]) || []
                }
                outputImages={
                  (sseData[currentDataIndex]?.outputImages as string[]) || []
                }
                chosenModel={sseData[currentDataIndex]?.chosenModel as string}
                onSwitchImageGroup={handleSwitchImageGroup}
              />
            </ReactFlowProvider>
          )}
        </div>
      </div>
    </div>
  );
}
